import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import confusion_matrix, classification_report

plt.rcParams['font.sans-serif'] = 'SimHei'
plt.rcParams['axes.unicode_minus'] = False

def evaluate_model():
    # 把训练好的模型加载出来
    model = tf.keras.models.load_model('mnist_cnn_model.h5')
    # 加载测试数据，这里在模型训练的时候已经处理好了，无需再次处理
    data = np.load('test_data.npz')
    test_images = data['images']
    test_labels = data['labels']
    # 加载训练历史
    history = np.load('cnn_train_history.npy', allow_pickle=True).item()

    # 使用测试集进行预测
    predictions = model.predict(test_images)
    predicted_labels = np.argmax(predictions, axis=1)
    true_labels = np.argmax(test_labels, axis=1)

    # 计算准确率
    accuracy = np.mean(predicted_labels == true_labels)
    print(f'Test Accuracy: {accuracy * 100:.2f}%')

    # 绘制损失曲线
    plt.figure(figsize=(10, 5))
    epochs = range(1, len(history['loss']) + 1)  # Epoch从1开始计数
    plt.plot(epochs, history['loss'], label='训练集损失函数')
    plt.plot(epochs, history['val_loss'], label='验证集损失函数')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title('Training and Validation Loss Curves')
    plt.legend()
    plt.savefig('loss_curve.png')  # 保存损失曲线图像
    plt.show()

    # 生成混淆矩阵并可视化
    cm = confusion_matrix(true_labels, predicted_labels)
    plt.figure(figsize=(10, 8))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
    plt.xlabel('预测标签')
    plt.ylabel('真实标签')
    plt.title('混淆矩阵')
    plt.savefig('confusion_matrix.png')  # 保存混淆矩阵图像
    plt.show()

    # 可视化部分测试样本及其预测结果
    num_samples = 10
    plt.figure(figsize=(15, 4))
    for i in range(num_samples):
        img = test_images[i].reshape(28, 28)  # 恢复图像形状
        plt.subplot(2, num_samples // 2, i + 1)  # 根据样本数调整布局
        plt.imshow(img, cmap='gray')
        plt.title(f'T:{true_labels[i]} P:{predicted_labels[i]}')
        plt.axis('off')
    plt.suptitle('预测样本结果示例')
    plt.savefig('sample_predictions.png')  # 保存样本图像


    misclassified_indices = np.where(predicted_labels != true_labels)[0]
    num_misclassified = len(misclassified_indices)
    if num_misclassified > 0:
        num_samples = min(10, num_misclassified)
        # 随机选择不重复的样本
        selected_indices = np.random.choice(misclassified_indices, num_samples, replace=False)

        # 创建自适应布局
        n_cols = 5
        n_rows = (num_samples + n_cols - 1) // n_cols
        fig, axes = plt.subplots(n_rows, n_cols, figsize=(15, 4), squeeze=False)
        axes = axes.flatten()

        # 绘制每个误判样本
        for i in range(num_samples):
            img = test_images[selected_indices[i]].reshape(28, 28)
            axes[i].imshow(img, cmap='gray')
            axes[i].set_title(f'T:{true_labels[selected_indices[i]]} P:{predicted_labels[selected_indices[i]]}',
                              color='red')
            axes[i].axis('off')

        # 隐藏空白子图
        for j in range(num_samples, len(axes)):
            axes[j].axis('off')

        plt.suptitle('错误分类样本示例')
        plt.tight_layout()
        plt.savefig('misclassified_samples.png', bbox_inches='tight')
        plt.show()
    else:
        print("所有的样本都被正确预测了!")
    print('Classification Report:')
    print(classification_report(true_labels, predicted_labels))
if __name__ == '__main__':
        evaluate_model()